The K weakest rows [OrderedDict, Quick Select]¶
Time: O(MxN); Space: O(K); easy
Given a m * n matrix mat of ones (representing soldiers) and zeros (representing civilians), return the indexes of the k weakest rows in the matrix ordered from the weakest to the strongest.
A row i is weaker than row j, if the number of soldiers in row i is less than the number of soldiers in row j, or they have the same number of soldiers but i is less than j. Soldiers are always stand in the frontier of a row, that is, always ones may appear first and then zeros.
Example 1:
Input: mat =
[[1,1,0,0,0], [1,1,1,1,0], [1,0,0,0,0], [1,1,0,0,0], [1,1,1,1,1]], k = 3 Output: [2, 0, 3]
Explanation:
The number of soldiers for each row is: row 0 -> 2 row 1 -> 4 row 2 -> 1 row 3 -> 2 row 4 -> 5
Rows ordered from the weakest to the strongest are [2,0,3,1,4]
Example 2:
Input: mat =
[[1,0,0,0], [1,1,1,1], [1,0,0,0], [1,0,0,0]], k = 2 Output: [0, 2]
Explanation:
The number of soldiers for each row is: row 0 -> 1 row 1 -> 4 row 2 -> 1 row 3 -> 1
Rows ordered from the weakest to the strongest are [0,2,3,1]
Notes:
len(mat) = m
len(mat[i]) = n
2 <= n, m <= 100
1 <= k <= m
matrix[i][j] is either 0 or 1
Hints:
Sort the matrix row indexes by the number of soldiers and then row indexes.
[1]:
class Solution1(object):
"""
Time: O(m * n)
Space: O(k)
"""
def kWeakestRows(self, mat, k):
"""
:type mat: List[List[int]]
:type k: int
:rtype: List[int]
"""
result, lookup = [], set()
for j in range(len(mat[0])):
for i in range(len(mat)):
if mat[i][j] or i in lookup:
continue
lookup.add(i)
result.append(i)
if len(result) == k:
return result
for i in range(len(mat)):
if i in lookup:
continue
lookup.add(i)
result.append(i)
if len(result) == k:
break
return result
[2]:
s = Solution1()
mat = [
[1,1,0,0,0],
[1,1,1,1,0],
[1,0,0,0,0],
[1,1,0,0,0],
[1,1,1,1,1]
]
k = 3
assert s.kWeakestRows(mat, k) == [2, 0, 3]
mat =[
[1,0,0,0],
[1,1,1,1],
[1,0,0,0],
[1,0,0,0]
]
k = 2
assert s.kWeakestRows(mat, k) == [0,2]
[15]:
import collections
class Solution2(object):
"""
Time: O(m * n)
Space: O(k)
"""
def kWeakestRows(self, mat, k):
"""
:type mat: List[List[int]]
:type k: int
:rtype: List[int]
"""
lookup = collections.OrderedDict()
for j in range(len(mat[0])):
for i in range(len(mat)):
if mat[i][j] or i in lookup:
continue
lookup[i] = True
if len(lookup) == k:
return list(lookup.keys())
for i in range(len(mat)):
if i in lookup:
continue
lookup[i] = True
if len(lookup) == k:
break
return list(lookup.keys())
[17]:
s = Solution2()
mat = [
[1,1,0,0,0],
[1,1,1,1,0],
[1,0,0,0,0],
[1,1,0,0,0],
[1,1,1,1,1]
]
k = 3
assert s.kWeakestRows(mat, k) == [2, 0, 3]
mat =[
[1,0,0,0],
[1,1,1,1],
[1,0,0,0],
[1,0,0,0]
]
k = 2
assert s.kWeakestRows(mat, k) == [0,2]
[18]:
import random
class Solution3(object):
"""
Time: O(m * n + klogk)
Space: O(m)
"""
def kWeakestRows(self, mat, k):
"""
:type mat: List[List[int]]
:type k: int
:rtype: List[int]
"""
def nth_element(nums, n, compare=lambda a, b: a < b):
def partition_around_pivot(left, right, pivot_idx, nums, compare):
new_pivot_idx = left
nums[pivot_idx], nums[right] = nums[right], nums[pivot_idx]
for i in range(left, right):
if compare(nums[i], nums[right]):
nums[i], nums[new_pivot_idx] = nums[new_pivot_idx], nums[i]
new_pivot_idx += 1
nums[right], nums[new_pivot_idx] = nums[new_pivot_idx], nums[right]
return new_pivot_idx
left, right = 0, len(nums) - 1
while left <= right:
pivot_idx = random.randint(left, right)
new_pivot_idx = partition_around_pivot(left, right, pivot_idx, nums, compare)
if new_pivot_idx == n:
return
elif new_pivot_idx > n:
right = new_pivot_idx - 1
else: # new_pivot_idx < n
left = new_pivot_idx + 1
nums = [(sum(mat[i]), i) for i in range(len(mat))]
nth_element(nums, k)
return list(map(lambda x: x[1], sorted(nums[:k])))
[19]:
s = Solution3()
mat = [
[1,1,0,0,0],
[1,1,1,1,0],
[1,0,0,0,0],
[1,1,0,0,0],
[1,1,1,1,1]
]
k = 3
assert s.kWeakestRows(mat, k) == [2, 0, 3]
mat =[
[1,0,0,0],
[1,1,1,1],
[1,0,0,0],
[1,0,0,0]
]
k = 2
assert s.kWeakestRows(mat, k) == [0,2]